//#define OUTPUT_RESULT

#include"../interchiplet_app/interchiplet_app.h"

#include<iostream>
#include<numeric>
#include<random>
#include<array>
#include<thread>
#include<cstdlib>
#include<algorithm>
#include<functional>
#include<cstring>
#include<cstdio>
#include<unistd.h>
#include<ctime>
#include<chrono>

using namespace std;

//矩阵大小
const size_t N=300;
constexpr int64_t CHIPLET_COUNT=8;
const size_t THREAD_COUNT=24;
const size_t LOOP=50000;
typedef array<array<double,N>,N> matrix_t;

int seed=0;
int64_t chipletId=0;
matrix_t*a,*b,*c;

nsInterchiplet::syscall_return_t neo_sendMsg(int64_t localPort,int64_t datum)
{
    bool b=true;
    thread th([&](){
        volatile int i=0;
        while(b)i++;
    });
    nsInterchiplet::syscall_return_t y=nsInterchiplet::sendMessage(localPort,datum);
    b=false;
    th.join();
    return y;
}
nsInterchiplet::syscall_return_t neo_recvMsg(int64_t localPort,int64_t*desPtr)
{
    bool b=true;
    thread th([&](){
        volatile int i=0;
        while(b)i++;
    });
    nsInterchiplet::syscall_return_t y=nsInterchiplet::receiveMessage(localPort,desPtr);
    b=false;
    th.join();
    return y;
}

void initMtrx(mt19937&generator,uniform_real_distribution<>&urd,matrix_t&mtrx)
{
    for(auto&line:mtrx){
        generate(line.begin(),line.end(),[&](){
            return urd(generator);
        });
    }
}
void multiply(matrix_t&mtrxA,matrix_t&mtrxB,matrix_t&mtrxC)
{
    thread*th[THREAD_COUNT];
    int i;
    for(i=0;i<THREAD_COUNT;++i){
        th[i]=new thread([&](int idx){
            for(size_t k=idx;k<N;k+=THREAD_COUNT)
                for(int64_t j=chipletId;j<N;j+=CHIPLET_COUNT)
                    mtrxC[j][k]=inner_product(mtrxA[j].begin(),mtrxA[j].end(),mtrxB[k].begin(),0.0);
        },i);
    }
    for(i=0;i<THREAD_COUNT;++i)th[i]->join();
    for(i=0;i<THREAD_COUNT;++i)delete th[i];
}
void sendResults()
{
    using namespace nsInterchiplet;
    int64_t port=chipletId;
    openPort(port,0,port);
    int64_t t;
    int i;
    for(i=0;i<N;++i){
        memcpy(&t,&(*c)[chipletId][i],sizeof(double));
        sendMessage(port,t);
        //neo_sendMsg(port,t);
    }
    closePort(port);
}
void sendResults2()
{
    thread*th[N];
    int i;
    for(i=0;i<N;++i){
        th[i]=new thread([&](int idx){
            int64_t port=chipletId*N+idx;
            //int64_t t;
            nsInterchiplet::openPort(port,0,port);
            nsInterchiplet::sendMessage(port,*(int64_t*)&(*c)[chipletId][idx]);
            //neo_sendMsg(port,*(int64_t*)&(*c)[chipletId][idx]);
            //cout<<"send msg: "<<nsInterchiplet::sendMessage(port,*(int64_t*)&(*c)[chipletId][idx])<<endl;
            nsInterchiplet::closePort(port);
        },i);
    }
    for(i=0;i<N;++i){
        th[i]->join();
        delete th[i];
    }
}
void sendResults3()
{
    bool b=true;
    thread*th[THREAD_COUNT];
    size_t i;
    for(i=0;i<THREAD_COUNT;++i){
        th[i]=new thread([&](size_t idx){
            for(size_t k=idx;k<N;k+=THREAD_COUNT){
                int64_t port=chipletId*N+k;
                nsInterchiplet::openPort(port,0,port);
                for(int64_t j=chipletId;j<N;j+=CHIPLET_COUNT){
                    nsInterchiplet::sendMessage(port,*(int64_t*)&(*c)[j][k]);
                    //neo_sendMsg(port,*(int64_t*)&(*c)[j][k]);
                }
                nsInterchiplet::closePort(port);
            }
        },i);
    }
    //for(i=0;i<LOOP;++i);
    thread runningThread([&](){
        volatile int r=0;
        while(b){
            r++;
            this_thread::yield();
        }
    });
    for(auto&t:th){
        t->join();
        delete t;
    }
    b=false;
    runningThread.join();
}
void open_port()
{
    for(int64_t i=1;i<N;++i)nsInterchiplet::openPort(i,i,i);
}
void open_port2()
{
    int64_t i,port_max=N*N;
    for(i=N;i<port_max;++i)nsInterchiplet::openPort(i,i/N,i);
}
void recvResults()
{
    using namespace nsInterchiplet;
    int64_t i;
    //for(i=1;i<N;++i)openPort(i,i,i);
    for(i=1;i<N;++i){
        for(int j=0;j<N;++j){
            //receiveMessage(i,(int64_t*)&(*c)[i][j]);
            neo_recvMsg(i,(int64_t*)&(*c)[i][j]);
        }
    }
    for(i=1;i<N;++i)closePort(i);
}
void recvResults2()
{
    thread*th[N];
    int i;
    for(i=1;i<N;++i){
        th[i]=new thread([&](int idx){
            int64_t port_min=N*idx;
            int64_t port_max=port_min+N;
            int64_t j;
            int64_t t;
            //if(i>1)th[idx-1]->join();
            for(j=port_min;j<port_max;++j){
                /* cout<<"Thread "<<idx<<" open port "<<j<<endl;
                nsInterchiplet::openPort(j,idx,j); */
                //cout<<"Thread "<<idx<<" receiving from port "<<j<<endl;
                nsInterchiplet::receiveMessage(j,/* (int64_t*)&(*c)[idx][j-port_min] */&t);
                //neo_recvMsg(j,&t);
                //t=syscall(501,j,nullptr);
                //cout<<"Thread "<<idx<<" copying\n";
                memcpy(&(*c)[idx][j-port_min],&t,sizeof(double));
                //cout<<"Thread "<<idx<<" finish receiving\n";
                nsInterchiplet::closePort(j);
            }
        },i);
    }
    for(i=0;i<10000;++i){
        //do nothing
    }
    for(i=1;i<N;++i){
        th[i]->join();
        delete th[i];
    }
}
void recvResults2b()
{
    int64_t i,port_max=N*N;
    for(i=N;i<port_max;++i){
        nsInterchiplet::receiveMessage(i,(int64_t*)&(*c)[i/N][i%N]);
        //neo_recvMsg(i,(int64_t*)&(*c)[i/N][i%N]);
        nsInterchiplet::closePort(i);
    }
}
void recvResults2c()
{
    thread th([&](){
        this_thread::sleep_for(1ms);
        //recvResults2b();
    });
    recvResults2b();
    th.join();
}
void open_port3()
{
    for(int64_t i=1;i<CHIPLET_COUNT;++i){
        int64_t port_base=i*N;
        for(int64_t j=0;j<N;++j){
            int64_t port=port_base+j;
            nsInterchiplet::openPort(port,i,port);
        }
    }
}
void recvResults3()
{
    using namespace nsInterchiplet;
    bool b=true;
    thread*th[THREAD_COUNT];
    for(size_t i=0;i<THREAD_COUNT;++i){
        th[i]=new thread([&](size_t idx){
            for(size_t l=idx;l<N;l+=THREAD_COUNT){
                for(int64_t j=1;j<CHIPLET_COUNT;++j){
                    int64_t port=j*N+l;
                    for(int64_t k=j;k<N;k+=CHIPLET_COUNT){
                        receiveMessage(port,(int64_t*)&(*c)[k][l]);
                        //neo_recvMsg(port,(int64_t*)&(*c)[k][l]);
                    }
                    closePort(port);
                }
            }
        },i);
    }
    thread runningThread([&](){
        volatile int r=0;
        while(b){
            r++;
            this_thread::yield();
        }
    });
    for(auto&t:th){
        t->join();
        delete t;
    }
    b=false;
    runningThread.join();
}

//随机数种子
int main(int argc,char**argv)
{
    if(argc>2){
        cout<<"Wrong number of parameters\n";
        return 1;
    }
    cout<<"start\n";
    //cout<<"start at "<<clock()<<"\n";
    if(argc==2)seed=atoi(argv[1]);
    //#if CHIPLET_COUNT > 1
    chipletId=nsInterchiplet::getLocalAddress();
    //#endif
    if(chipletId<0){
        cout<<"Interchiplet error\n";
        return 2;
    }
    mt19937 gen(seed);
    uniform_real_distribution<>dis(0,1);
    a=new matrix_t;
    b=new matrix_t;
    c=new matrix_t;
    if(chipletId==0)open_port3();
    initMtrx(gen,dis,*a);
    initMtrx(gen,dis,*b);
    cout<<"computing...\n";
    //cout<<"computing at "<<clock()<<"...\n";
    multiply(*a,*b,*c);
    if(chipletId==0){
        cout<<"finished compution, waiting for other results...\n";
        //cout<<"Finished compution, waiting for other results at "<<clock()<<"...\n";
        recvResults3();
        #ifdef OUTPUT_RESULT
        cout<<"Chiplet "<<chipletId<<":\n";
        for(auto&vec:*c){
            for(auto&value:vec)cout<<value<<' ';
            cout<<endl;
        }
        #endif
    }
    else{
        cout<<"sending result...\n";
        //cout<<"sending result at "<<clock()<<"...\n";
        sendResults3();
        #ifdef OUTPUT_RESULT
        cout<<"Chiplet "<<chipletId<<":\n";
        for(int64_t i=chipletId;i<N;i+=CHIPLET_COUNT){
            for(auto&v:(*c)[i])cout<<v<<' ';
            cout<<endl;
        }
        #endif
    }
    delete[]a,b,c;
    cout<<"the end\n";
    //cout<<"The end at "<<clock()<<"\n";
    return 0;
}
